NEP_UDF_GEOTIFF_BIGTIFF = 0x0040     # NC_UDF0
NEP_UDF_GEOTIFF_STANDARD = 0x0080    # NC_UDF1
NEP_UDF_GRIB2 = 0x10000              # NC_UDF2
NEP_UDF_FITS = 0x80000               # NC_UDF3
NEP_UDF_CDF = 0x100000               # NC_UDF4
NEP_UDF_PDS4 = 0x200000              # NC_UDF5
NEP_UDF_DICOM = 0x400000             # NC_UDF6
NEP_UDF_PDB = 0x800000               # NC_UDF7
NEP_UDF_MMCIF = 0x1000000            # NC_UDF8
NEP_UDF_NEXTCDF4 = 0x2000000         # NC_UDF9

# User-defined format (UDF) mode flags, used with the extra_mode_flag
# Dataset kwarg. netcdf-c has used two encodings for these:
#  - netcdf-c 4.10.x: one mode bit per slot, slots 0-9 ("bitmask").
#  - later netcdf-c: NC_UDF(n) = NC_UDF_FLAG | (n << NC_UDF_NUM_SHIFT),
#    slots 0-63 ("slotfield").
# Only slot 0 (0x0040) has the same value in both, so the encoding is
# detected from the loaded library instead of being hard-coded.
_UDF_FLAG = 0x0040
_UDF_NUM_SHIFT = 19
_UDF_BITMASK_FLAGS = (0x0040, 0x0080, 0x10000, 0x80000, 0x100000,
                      0x200000, 0x400000, 0x800000, 0x1000000, 0x2000000)

cdef _detect_udf_encoding():
    if _nclibversion < Version("4.10.0"):
        return None, 0
    # 0x0080 is NC_UDF1 in the bitmask encoding. In the slotfield encoding
    # it lacks NC_UDF_FLAG, so nc_inq_user_format rejects it.
    if nc_inq_user_format(0x0080, NULL, NULL) == NC_NOERR:
        return "bitmask", len(_UDF_BITMASK_FLAGS)
    return "slotfield", 64

_udf_encoding, _udf_max_slots = _detect_udf_encoding()


def udf_mode(slot):
    """
**`udf_mode(slot)`**

Return the mode flag that selects user-defined format (UDF) slot `slot`
in the netcdf-c library in use, for the `extra_mode_flag` kwarg of
`Dataset`. The numeric value depends on the netcdf-c version, so use this
function (or the `NEP_UDF_*` constants) instead of hard-coding integers.
Valid slots are 0-9 for netcdf-c 4.10.x and 0-63 for later versions.
Requires netcdf-c >= 4.10.0."""
    if _udf_encoding is None:
        raise NetCDF4MissingFeatureException("user-defined formats", "4.10.0")
    if not isinstance(slot, int) or isinstance(slot, bool):
        raise TypeError(f"UDF slot must be an int, got {slot!r}")
    if not 0 <= slot < _udf_max_slots:
        raise ValueError(f"UDF slot must be in range 0-{_udf_max_slots - 1}, got {slot}")
    if _udf_encoding == "bitmask":
        return _UDF_BITMASK_FLAGS[slot]
    return _UDF_FLAG | (slot << _UDF_NUM_SHIFT)


def _udf_registered(slot):
    """True if a dispatch layer is registered with netcdf-c for UDF slot `slot`."""
    cdef NC_Dispatch *dispatch = NULL
    cdef int mode = udf_mode(slot)
    cdef int ierr
    with nogil:
        ierr = nc_initialize()  # loads UDF plugins listed in .ncrc; no-op if already done
    _ensure_nc_success(ierr)
    with nogil:
        ierr = nc_inq_user_format(mode, &dispatch, NULL)
    _ensure_nc_success(ierr)
    return dispatch != NULL


# All valid UDF modes for the loaded library (empty if UDFs unsupported).
_udf_modes = frozenset(udf_mode(n) for n in range(_udf_max_slots))

# Mode flags for the user-defined formats provided by the NetCDF Expansion
# Pack (NEP). Names and slot numbers match those in nep.h. If the library has
# no UDF support the 4.10.x values are used, so the names always exist;
# passing them to Dataset then raises NetCDF4MissingFeatureException.
def _nep_udf(slot):
    return udf_mode(slot) if _udf_encoding is not None else _UDF_BITMASK_FLAGS[slot]

NEP_UDF_GEOTIFF_BIGTIFF = _nep_udf(0)
NEP_UDF_GEOTIFF_STANDARD = _nep_udf(1)
NEP_UDF_GRIB2 = _nep_udf(2)
NEP_UDF_FITS = _nep_udf(3)
NEP_UDF_CDF = _nep_udf(4)
NEP_UDF_PDS4 = _nep_udf(5)
NEP_UDF_DICOM = _nep_udf(6)
NEP_UDF_PDB = _nep_udf(7)
NEP_UDF_MMCIF = _nep_udf(8)
NEP_UDF_NEXTCDF4 = _nep_udf(9)
del _nep_udf

        # extra mode flag for user-defined formats, only used for plain
        # read-only nc_open (cmode == NC_NOWRITE). Need to check read_mode
        # too, since NC_CLOBBER (mode='w') is also zero.
        if extra_mode_flag is not None:
            if _udf_encoding is None:
                raise NetCDF4MissingFeatureException("extra_mode_flag", "4.10.0")
            if not read_mode or cmode != NC_NOWRITE or memory is not None or parallel:
                raise ValueError("extra_mode_flag only allowed with mode='r', "
                                 "and not with diskless, memory or parallel")
            if extra_mode_flag not in _udf_modes:
                raise ValueError("extra_mode_flag must be one of the NEP_UDF_* "
                                 "constants or a value returned by udf_mode(), "
                                 f"got {extra_mode_flag!r}")

__all__ = [
    'Dataset', 'Variable', 'Dimension', 'Group', 'MFDataset', 'MFTime', 'CompoundType',
    'VLType', 'date2num', 'num2date', 'date2index', 'stringtochar', 'chartostring',
    'stringtoarr', 'getlibversion', 'EnumType', 'get_chunk_cache', 'set_chunk_cache',
    'set_alignment', 'get_alignment', 'rc_get', 'rc_set', 'udf_mode',
]

NEP_UDF_GEOTIFF_BIGTIFF: int
NEP_UDF_GEOTIFF_STANDARD: int
NEP_UDF_GRIB2: int
NEP_UDF_FITS: int
NEP_UDF_CDF: int
NEP_UDF_PDS4: int
NEP_UDF_DICOM: int
NEP_UDF_PDB: int
NEP_UDF_MMCIF: int
NEP_UDF_NEXTCDF4: int

def udf_mode(slot: int) -> int: ...

import unittest, os, tempfile
from numpy.testing import assert_array_almost_equal
import netCDF4
from netCDF4._netCDF4 import _udf_encoding, _udf_max_slots, _udf_registered

# test extra_mode_flag Dataset kwarg (issue #1490), used to select a
# user-defined format (UDF) dispatch layer for files with no magic number.
# The mmCIF read test only runs if the NetCDF Expansion Pack (NEP) mmCIF
# reader is registered with netcdf-c (e.g. via NCRCENV_RC pointing to NEP's .ncrc).

NC_ENOTBUILT = -128

MMCIF = """\
data_TINY
#
_entry.id TINY
#
loop_
_atom_site.group_PDB
_atom_site.id
_atom_site.type_symbol
_atom_site.label_atom_id
_atom_site.label_alt_id
_atom_site.label_comp_id
_atom_site.label_asym_id
_atom_site.label_entity_id
_atom_site.label_seq_id
_atom_site.pdbx_PDB_ins_code
_atom_site.Cartn_x
_atom_site.Cartn_y
_atom_site.Cartn_z
_atom_site.occupancy
_atom_site.B_iso_or_equiv
_atom_site.auth_seq_id
_atom_site.auth_asym_id
_atom_site.pdbx_PDB_model_num
ATOM 1 N N . VAL A 1 1 ? 6.204 16.869 4.854 1.00 49.05 1 A 1
ATOM 2 C CA . VAL A 1 1 ? 6.913 17.759 4.607 1.00 43.14 1 A 1
ATOM 3 C C . VAL A 1 1 ? 8.504 17.378 4.797 1.00 24.80 1 A 1
#
"""
CARTN_X = [6.204, 6.913, 8.504]

# slot numbers from NEP's nep.h
NEP_SLOTS = dict(GEOTIFF_BIGTIFF=0, GEOTIFF_STANDARD=1, GRIB2=2, FITS=3, CDF=4,
                 PDS4=5, DICOM=6, PDB=7, MMCIF=8, NEXTCDF4=9)
BITMASK = (0x0040, 0x0080, 0x10000, 0x80000, 0x100000,
           0x200000, 0x400000, 0x800000, 0x1000000, 0x2000000)

has_udf = _udf_encoding is not None
has_nep_mmcif = has_udf and _udf_registered(NEP_SLOTS['MMCIF'])


def _write(text, suffix):
    f = tempfile.NamedTemporaryFile(suffix=suffix, delete=False, mode='w')
    f.write(text)
    f.close()
    return f.name


@unittest.skipIf(has_udf, "only for netcdf-c < 4.10.0")
class NoUDFSupportTestCase(unittest.TestCase):

    def runTest(self):
        """extra_mode_flag and udf_mode need netcdf-c >= 4.10.0"""
        with self.assertRaises(netCDF4.NetCDF4MissingFeatureException):
            netCDF4.udf_mode(0)
        with self.assertRaises(netCDF4.NetCDF4MissingFeatureException):
            netCDF4.Dataset('unused.nc', extra_mode_flag=netCDF4.NEP_UDF_MMCIF)


@unittest.skipUnless(has_udf, "extra_mode_flag requires netcdf-c >= 4.10.0")
class UDFModeTestCase(unittest.TestCase):

    def runTest(self):
        """udf_mode values match the encoding used by the loaded netcdf-c"""
        self.assertIn(_udf_encoding, ("bitmask", "slotfield"))
        self.assertEqual(_udf_max_slots, 10 if _udf_encoding == "bitmask" else 64)
        for slot in range(_udf_max_slots):
            mode = netCDF4.udf_mode(slot)
            if _udf_encoding == "bitmask":
                self.assertEqual(mode, BITMASK[slot])
            else:
                self.assertEqual(mode, 0x0040 | (slot << 19))
        # NC_UDF0 is the same in both encodings
        self.assertEqual(netCDF4.udf_mode(0), 0x0040)
        modes = [netCDF4.udf_mode(s) for s in range(_udf_max_slots)]
        self.assertEqual(len(set(modes)), len(modes))
        # NEP constants map to the nep.h slot numbers
        for name, slot in NEP_SLOTS.items():
            self.assertEqual(getattr(netCDF4, 'NEP_UDF_' + name), netCDF4.udf_mode(slot))
        for bad in (-1, _udf_max_slots):
            with self.assertRaises(ValueError):
                netCDF4.udf_mode(bad)
        for bad in (1.0, '1', True):
            with self.assertRaises(TypeError):
                netCDF4.udf_mode(bad)


@unittest.skipUnless(has_udf, "extra_mode_flag requires netcdf-c >= 4.10.0")
class ExtraModeFlagArgsTestCase(unittest.TestCase):

    def setUp(self):
        self.file = tempfile.NamedTemporaryFile(suffix='.nc', delete=False).name
        netCDF4.Dataset(self.file, 'w').close()

    def tearDown(self):
        os.remove(self.file)

    def runTest(self):
        """only allowed for plain read-only opens with a valid UDF mode"""
        flag = netCDF4.NEP_UDF_MMCIF
        for kwargs in (dict(mode='w'), dict(mode='a'), dict(mode='r+'),
                       dict(mode='rs'), dict(mode='r', diskless=True)):
            with self.assertRaises(ValueError, msg=str(kwargs)):
                netCDF4.Dataset(self.file, extra_mode_flag=flag, **kwargs)
        with open(self.file, 'rb') as f:
            buf = f.read()
        with self.assertRaises(ValueError):
            netCDF4.Dataset('inmemory.nc', memory=buf, extra_mode_flag=flag)
        # not a UDF mode (NC_WRITE)
        with self.assertRaises(ValueError):
            netCDF4.Dataset(self.file, extra_mode_flag=0x0001)
        # a UDF mode from the *other* encoding is rejected
        other = (0x0040 | (1 << 19)) if _udf_encoding == "bitmask" else BITMASK[1]
        with self.assertRaises(ValueError):
            netCDF4.Dataset(self.file, extra_mode_flag=other)
        # extra_mode_flag=None is the default
        netCDF4.Dataset(self.file, extra_mode_flag=None).close()


@unittest.skipUnless(has_udf, "extra_mode_flag requires netcdf-c >= 4.10.0")
class UDFModeFlagReachesNcOpenTestCase(unittest.TestCase):

    def setUp(self):
        self.file = tempfile.NamedTemporaryFile(suffix='.nc', delete=False).name
        netCDF4.Dataset(self.file, 'w').close()
        self.unregistered = [s for s in range(_udf_max_slots) if not _udf_registered(s)]

    def tearDown(self):
        os.remove(self.file)

    def runTest(self):
        """netcdf-c sees the flag as a UDF selection (fails with NC_ENOTBUILT
        for an unregistered slot instead of opening the file as netCDF)"""
        if not self.unregistered:
            self.skipTest("all UDF slots are registered")
        for slot in self.unregistered:
            with self.assertRaises(OSError, msg=f"slot {slot}") as cm:
                netCDF4.Dataset(self.file, extra_mode_flag=netCDF4.udf_mode(slot))
            self.assertEqual(cm.exception.errno, NC_ENOTBUILT, msg=f"slot {slot}")


@unittest.skipUnless(has_nep_mmcif, "NEP mmCIF reader not registered")
class ExtraModeFlagNEPTestCase(unittest.TestCase):

    def setUp(self):
        # without the leading data_ magic number netcdf-c can't detect the format
        self.file = _write("# no magic number\n" + MMCIF, '.cif')

    def tearDown(self):
        os.remove(self.file)

    def runTest(self):
        """open mmCIF file with no magic number using NEP_UDF_MMCIF"""
        with self.assertRaises(OSError):
            netCDF4.Dataset(self.file)
        with netCDF4.Dataset(self.file, extra_mode_flag=netCDF4.NEP_UDF_MMCIF) as ds:
            assert ds.entry_id == 'TINY'
            assert len(ds.dimensions['atom']) == 3
            assert_array_almost_equal(ds['atom_site_Cartn_x'][0], CARTN_X)


if __name__ == '__main__':
    unittest.main()

from netCDF4._netCDF4 import _udf_encoding, _udf_registered

has_nep_grib2 = _udf_encoding is not None and _udf_registered(2)  # NEP_UDF_GRIB2 slot